import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False
np.random.seed(1986)


# 转换数据集
def create_dataset(dataset, look_back=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - look_back):
        a = dataset[i:(i + look_back), 0]
        dataX.append(a)
        dataY.append(dataset[i + look_back, 0])
    return np.array(dataX), np.array(dataY)


# 导入数据
data_source = pd.read_excel("airline.xlsx")
raw_dataset = data_source[['数量']].values

# 数据归一化,[0,1]
scaler = MinMaxScaler(feature_range=(0, 1))
dataset = scaler.fit_transform(raw_dataset)

# 按比例拆分训练集和测试集(连续)
train_ratio = 0.8
train_size = int(len(dataset) * train_ratio)
train = dataset[0:train_size, :]
test = dataset[train_size:len(dataset), :]

# 数据集按转look_back大小窗口转换为输入X和输出Y
look_back = 3
train_x, train_y = create_dataset(train, look_back)
test_x, test_y = create_dataset(test, look_back)

# 输入数据reshape
train_x = np.reshape(train_x, (train_x.shape[0], 1, train_x.shape[1]))
test_x = np.reshape(test_x, (test_x.shape[0], 1, test_x.shape[1]))

# 创建模型
model = Sequential()
model.add(LSTM(units=8, input_shape=(1, look_back)))
model.add(Dense(1))
model.compile(loss='mean_squared_error', optimizer='adam')

# 打印模型信息
model.summary()

# 训练模型
history = model.fit(train_x, train_y, epochs=10, batch_size=1, verbose=0)

# 绘制训练过程(损失函数)
plt.plot(history.history["loss"])
plt.title("模型损失函数")
plt.xlabel("epoch")
plt.ylabel("loss_value")
plt.show()

# 训练集
train_predict = model.predict(train_x)
train_predict = scaler.inverse_transform(train_predict)
train_y = scaler.inverse_transform(train_y.reshape(-1, 1))

train_score = math.sqrt(mean_squared_error(train_predict, train_y))
print('训练集 RMSE:%.2f' % (train_score))

# 绘图
plt.title("训练集")
plt.plot(train_y)
plt.plot(train_predict)
plt.legend(["真实值", "测试值"])
plt.show()

# 测试集
test_predict = model.predict(test_x)
test_predict = scaler.inverse_transform(test_predict)
test_y = scaler.inverse_transform(test_y.reshape(-1, 1))

# 绘图
plt.title("测试集")
plt.plot(test_y)
plt.plot(test_predict)
plt.legend(["真实值", "测试值"])
plt.show()

test_score = math.sqrt(mean_squared_error(test_predict, test_y))
print('测试集RMSE:%.2f' % (test_score))

# 预测下一个周期的值
input_x = dataset[-1 * look_back:, :]
input_x = np.reshape(input_x, (input_x.shape[1], 1, input_x.shape[0]))
output_y = model.predict(input_x)
output_y = scaler.inverse_transform(output_y.reshape(-1, 1))
output_y = output_y[0][0]
print("预测下一周期的值:%f" % (output_y))

